import base64
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.backends import default_backend
import json

from pyexpat.errors import messages

from app.config.settings import settings


class RSACipher:
    def __init__(self, private_key_path=None, public_key_path=None):
        if private_key_path and public_key_path:
            self.private_key = self.load_private_key(private_key_path)
            self.public_key = self.load_public_key(public_key_path)
        else:
            self.private_key, self.public_key = self.generate_key_pair()

    @staticmethod
    def generate_key_pair():
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        public_key = private_key.public_key()
        return private_key, public_key

    @staticmethod
    def load_private_key(path):
        with open(path, "rb") as key_file:
            private_key = serialization.load_pem_private_key(
                key_file.read(),
                password=None,
                backend=default_backend()
            )
        return private_key

    @staticmethod
    def load_public_key(path):
        with open(path, "rb") as key_file:
            public_key = serialization.load_pem_public_key(
                key_file.read(),
                backend=default_backend()
            )
        return public_key

    def encrypt(self, plaintext):
        if isinstance(plaintext, (dict, list)):
            plaintext = json.dumps(plaintext)
        if isinstance(plaintext, str):
            plaintext = plaintext.encode('utf-8')

        # 计算每次加密的最大长度
        max_length = (self.public_key.key_size // 8) - 2 * hashes.SHA256.digest_size - 2
        ciphertext_parts = []

        # 分段加密
        for i in range(0, len(plaintext), max_length):
            part = plaintext[i:i + max_length]
            ciphertext = self.public_key.encrypt(
                part,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            ciphertext_parts.append(ciphertext)

        # 合并所有分段的密文并进行 Base64 编码
        combined_ciphertext = b''.join(ciphertext_parts)
        return base64.b64encode(combined_ciphertext).decode('utf-8')

    def decrypt(self, ciphertext):
        if isinstance(ciphertext, str):
            ciphertext = base64.b64decode(ciphertext)

        # 计算每次解密的最大长度
        max_length = self.private_key.key_size // 8
        plaintext_parts = []

        # 分段解密
        for i in range(0, len(ciphertext), max_length):
            part = ciphertext[i:i + max_length]
            plaintext = self.private_key.decrypt(
                part,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            plaintext_parts.append(plaintext)

        # 合并所有分段的明文
        decrypted_text = b''.join(plaintext_parts).decode('utf-8')
        try:
            return json.loads(decrypted_text)
        except json.JSONDecodeError:
            return decrypted_text

    def save_private_key(self, path):
        pem = self.private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        with open(path, "wb") as key_file:
            key_file.write(pem)

    def save_public_key(self, path):
        pem = self.public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        with open(path, "wb") as key_file:
            key_file.write(pem)



cipher = RSACipher(
    private_key_path=settings.RSA_PRIVATE_KEY_PATH,
    public_key_path=settings.RSA_PUBLIC_KEY_PATH
)

if __name__ == "__main__":
    cipher = RSACipher(
        private_key_path="./private_key.pem",
        public_key_path="./public_key.pem"
    )
    # message = {
    #     "name": "John",
    #     "age": 30,
    #     "city": "New York",
    #     "email": "EMAIL",
    #     "email": "john@example.com",
    #     "contact": "123-456-7890",
    #     "address": "123 Main St",
    #     "password": "PASSWORD",
    #     "password": "password123",
    #     "token": "TOKEN",
    #     "token": "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789",
    #     "key": "KEY",
    #     "key": "mysecretkey123",
    #     "data": "DATA",
    #     "data": "This is some sensitive data",
    #     "message": "MESSAGE",Received request
    #     "message": "Hello, world!",
    #     "content": "CONTENT",
    #     "content": "This is a sample content",
    # }
    message = [{'username': 'admin', 'password': '111111'}, {'username': 'root', 'password': '111111'}]
    ciphertext = cipher.encrypt(message)
    plaintext = cipher.decrypt(ciphertext)
    print(f"Original message: {message}")
    print(f"Ciphertext: {ciphertext}")
    print(f"Decrypted message: {plaintext}")
    # 对比密文和解密后的明文是否一致
    # if message == json.loads(plaintext):
    #     print("Encryption and decryption successful!")

    print(cipher.decrypt("d8uh6k8vMuT0tZBlA+W0bQM9qhT8V4cnw3G8ogHc7o1vqlfurcXP8tPNwC7SQnVuv9dWxtwlVBiL3H2q4p7imwmkkJVG/reRMkGNmyeW2eVC/P1oAcnlKUwBF96NCTEY2ki5ahM9Z1beRvzZNEmgSp/fbRdYYh+kDGrE+ARzjxMCrLPlQ+RhKUClSdUxj5ovddlUmP6xcOQBc4pHO8FOODzlAJIPortMDOqynkZvHPxjSNcXac0ktBTHYf3KxtFc6kzmUuzBTNqanTkzB07e4Ec2wv+GlwB8bN1IrUKBJPEn2M8Xjdop/Dv2mJwHDxdm1ggABYCpJ+lLXq5icUfE9Q=="))
    print(cipher.decrypt("GK/7GAmtZT4nQA/FYTRqqATtuaI5ESem7Joty8fale4RPBCu72ot1TI3kcyHaZdGK3hA2zSpJSEswgTIyLtfwNNhd0zde0bOKaV59sNMVDjn2GvX+w6rb2TZ95UD61fyHV8PS3dccCHJw7QRuv9O7/Rqr5Uim5kbJG+zufM776nFYQMn4SiO4FZTFRVLd3YvDLRHXSvVXm6em7HbgU3FptwyjF78RJ4dWHZWdYw8Pe8kv0wDR8oDaC2vyBs40q+OjrnB+YcLyM/A3wYmxhdhH0vftlhT18syiKMXZv44NrAwDLhcbWAstusApQv2F0soH5snNe0DHiUncDRaW+tlLg=="))